package ssh

import (
	"bufio"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os"
	"path/filepath"
	"time"

	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"
)

func getKeyFile(path string) (ssh.Signer, error) {
	homeDir, err := os.UserHomeDir()
	if err != nil {
		return nil, err
	}
	var file string
	if filepath.IsAbs(path) {
		file = path
	} else {
		file = filepath.Join(homeDir, path)
	}
	buf, err := ioutil.ReadFile(file)
	if err != nil {
		return nil, err
	}
	pubkey, err := ssh.ParsePrivateKey(buf)
	if err != nil {
		return nil, err
	}
	return pubkey, nil
}

// User field should be a name of user on remote server (ex. john in ssh john@example.com).
// Server field should be a remote machine address (ex. example.com in ssh john@example.com)
// Key is a path to private key on your local machine.
// Port is SSH server port on remote machine.
// Then ensure your Key begins from '/' (ex. /.ssh/id_rsa)
type Client struct {
	host   string
	port   string
	user   string
	client *ssh.Client

	password string
	keyPath  string
	callback ssh.HostKeyCallback
}

type Options func(c *Client)

func Password(pwd string) Options {
	return func(c *Client) {
		c.password = pwd
	}
}

func Callback(callback ssh.HostKeyCallback) Options {
	return func(c *Client) {
		c.callback = callback
	}
}

func KeyPath(keyPath string) Options {
	return func(c *Client) {
		c.keyPath = keyPath
	}
}

// NewClient 创建一刻ssh的客户端, 可以Run来执行命令, 可以Scp拷贝文件或者目录, 还可以Stream来获取流的命令输出
func NewClient(host string, port int, user string, options ...Options) (*Client, error) {
	c := &Client{
		host:     host,
		port:     fmt.Sprintf("%v", port),
		user:     user,
		callback: ssh.InsecureIgnoreHostKey(),
	}
	for _, option := range options {
		option(c)
	}
	err := c.connect()
	if err != nil {
		return nil, err
	}
	return c, nil
}

func (c *Client) Close() error {
	return c.client.Close()
}

func (c *Client) connect() error {
	var auth []ssh.AuthMethod
	if c.password != "" {
		auth = append(auth, ssh.Password(c.password))
	}
	if c.keyPath != "" {
		if pubkey, err := getKeyFile(c.keyPath); err == nil {
			auth = append(auth, ssh.PublicKeys(pubkey))
		}
	}
	if sshAgent, err := net.Dial("unix", os.Getenv("SSH_AUTH_SOCK")); err == nil {
		auth = append(auth, ssh.PublicKeysCallback(agent.NewClient(sshAgent).Signers))
		defer sshAgent.Close()
	}
	config := &ssh.ClientConfig{
		User:            c.user,
		Auth:            auth,
		HostKeyCallback: c.callback,
	}
	client, err := ssh.Dial("tcp", c.host+":"+c.port, config)
	if err != nil {
		return fmt.Errorf("connect(%v:%v)(user:%v,pwd:%v,key:%v)错误\n%w", c.host, c.port, c.user, c.password, c.keyPath, err)
	}
	c.client = client
	return nil
}

func (c *Client) newSession() (*ssh.Session, error) {
	if c.client == nil {
		err := c.connect()
		if err != nil {
			return nil, err
		}
	}
	session, err := c.client.NewSession()
	if err != nil {
		return nil, fmt.Errorf("newSession(user:%v,pwd:%v,key:%v)错误\n%w", c.user, c.password, c.keyPath, err)
	}
	return session, nil
}

func (c *Client) Scp(src, dst string) error {
	_args := []interface{}{src, dst}
	session, err := c.newSession()

	if err != nil {
		return err
	}
	defer session.Close()

	targetFile := filepath.Base(dst)

	srcFile, srcErr := os.Open(src)

	if srcErr != nil {
		return srcErr
	}

	srcStat, statErr := srcFile.Stat()

	if statErr != nil {
		return statErr
	}

	go func() {
		w, _ := session.StdinPipe()
		fmt.Fprintln(w, "C0644", srcStat.Size(), targetFile)
		if srcStat.Size() > 0 {
			io.Copy(w, srcFile)
			fmt.Fprint(w, "\x00")
			w.Close()
		} else {
			fmt.Fprint(w, "\x00")
			w.Close()
		}
	}()
	scpCommand := fmt.Sprintf("scp -tr %s", dst)
	if err := session.Run(scpCommand); err != nil {
		return fmt.Errorf("Scp%v\n%w", _args, err)
	}
	return nil
}

func (c *Client) Stream(command string, timeout int) (stdout chan string, stderr chan string, done chan bool, err error) {
	session, err := c.newSession()
	if err != nil {
		return stdout, stderr, done, err
	}
	outReader, err := session.StdoutPipe()
	if err != nil {
		return stdout, stderr, done, err
	}
	errReader, err := session.StderrPipe()
	if err != nil {
		return stdout, stderr, done, err
	}
	// combine outputs, create a line-by-line scanner
	stdoutReader := io.MultiReader(outReader)
	stderrReader := io.MultiReader(errReader)
	err = session.Start(command)
	stdoutScanner := bufio.NewScanner(stdoutReader)
	stderrScanner := bufio.NewScanner(stderrReader)
	// continuously send the command's output over the channel
	stdoutChan := make(chan string)
	stderrChan := make(chan string)
	done = make(chan bool)

	go func(stdoutScanner, stderrScanner *bufio.Scanner, stdoutChan, stderrChan chan string, done chan bool) {
		defer close(stdoutChan)
		defer close(stderrChan)
		defer close(done)

		timeoutChan := time.After(time.Duration(timeout) * time.Second)
		res := make(chan bool, 1)

		go func() {
			for stdoutScanner.Scan() {
				stdoutChan <- stdoutScanner.Text()
			}
			for stderrScanner.Scan() {
				stderrChan <- stderrScanner.Text()
			}
			// close all of our open resources
			res <- true
		}()

		select {
		case <-res:
			stdoutChan <- ""
			stderrChan <- ""
			done <- true
		case <-timeoutChan:
			stdoutChan <- ""
			stderrChan <- "Run Command Timeout!"
			done <- false
		}
		session.Close()
	}(stdoutScanner, stderrScanner, stdoutChan, stderrChan, done)
	return stdoutChan, stderrChan, done, err
}

func (c *Client) Run(command string, timeout ...int) (outStr string, errStr string, isTimeout bool, err error) {
	timeOutSecond := 5 * 60
	if len(timeout) > 0 {
		timeOutSecond = timeout[0]
	}
	stdoutChan, stderrChan, doneChan, err := c.Stream(command, timeOutSecond)
	if err != nil {
		return outStr, errStr, isTimeout, err
	}
	stillGoing := true
	for stillGoing {
		select {
		case isTimeout = <-doneChan:
			stillGoing = false
		case outline := <-stdoutChan:
			outStr += outline + "\n"
		case errline := <-stderrChan:
			errStr += errline + "\n"
		}
	}
	// return the concatenation of all signals from the output channel
	return outStr, errStr, isTimeout, err
}
